import torch
import gurobipy as gp
from gurobipy import GRB
import numpy as np
import csv

from prob import *
import verify_piecewise_utils
from NN import ann

# debug
#m.computeIIS()
#m.write("model1.ilp")

#landa = LANDA
#gama = GAMA
#x_var_num = x_var_num

# case_name = "eg2"
'''
    True : 不存在反例
    False: 存在反例
    
    0: X0 初始区域反例
    1: Xu 不安全区域反例
    2： X 蕴含条件反例
'''

# ————————————————————————————不安全区域的验证————————————————————————————
def MILP_opt_unsafeCon(Bx):
    cex_unsafe = []
    res_unsafe = []

    for i in range(len(sub_unsafe)):
        m = gp.Model()

        # 配置
        m.setParam('Outputflag', 0)
        m.setParam('NonConvex', 2)
        # m.setParam('InfUnbdInfo', 1)

        sub_unsafe_min = sub_unsafe[i][0]
        sub_unsafe_max = sub_unsafe[i][1]

        x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x0', lb=-np.inf, ub=np.inf)
        m.addConstr(x0[0] >= sub_unsafe_min[0])
        m.addConstr(x0[0] <= sub_unsafe_max[0])
        m.addConstr(x0[1] >= sub_unsafe_min[1])
        m.addConstr(x0[1] <= sub_unsafe_max[1])

        if sub_unsafe_shape[i] == 2:
            p0 = (sub_unsafe_min[0] + sub_unsafe_max[0]) / 2.0
            p1 = (sub_unsafe_min[1] + sub_unsafe_max[1]) / 2.0
            r = p0 - sub_unsafe_min[0]
            m.addConstr((x0[0]-p0)*(x0[0]-p0) + (x0[1]-p1)*(x0[1]-p1) <= r*r)

        y_layer = MILP_Encode_NN(m, x0, Bx)

        m.setObjective(y_layer, GRB.MINIMIZE)

        # 优化
        m.optimize()

        #m.computeIIS()
        #m.write("model1.ilp")
        counter_ex, min = solution_output(m)

        if min > 0.:
            # 不存在反例
            print(f"不安全区域{i}没有反例")
        else:
            # 存在反例
            print(f"不安全区域{i}反例是 : {counter_ex}")
            cex_unsafe.append(np.array(counter_ex))
            res_unsafe.append(min)
    print("cex_unsafe: ", cex_unsafe)
    print("res_unsafe: ", res_unsafe)
    return cex_unsafe, res_unsafe

# ————————————————————————————初始区域验证————————————————————————————
def MILP_opt_initCon(Bx):
    cex_init = []
    res_init = []

    for i in range(len(sub_init)):
        m = gp.Model()

        # 配置
        m.setParam('Outputflag', 0)
        m.setParam('NonConvex', 2)
        # m.setParam('InfUnbdInfo', 1)

        sub_init_min = sub_init[i][0]
        sub_init_max = sub_init[i][1]
        # Xi 区域约束
        x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x0', lb=-np.inf, ub=np.inf)
        m.addConstr(x0[0] >= sub_init_min[0])
        m.addConstr(x0[0] <= sub_init_max[0])
        m.addConstr(x0[1] >= sub_init_min[1])
        m.addConstr(x0[1] <= sub_init_max[1])

        if sub_init_shape[i] == 2:
            p0 = (sub_init_min[0] + sub_init_max[0]) / 2.0
            p1 = (sub_init_min[1] + sub_init_max[1]) / 2.0
            r = p0 - sub_init_min[0]
            m.addConstr((x0[0]-p0)*(x0[0]-p0) + (x0[1]-p1)*(x0[1]-p1) <= r*r)

        y_layer_for_x0 = MILP_Encode_NN(m, x0, Bx)

        m.setObjective(y_layer_for_x0, GRB.MAXIMIZE)

        m.optimize()

        counter_ex, maxi = solution_output(m)

        if maxi <= 0:
            # 不存在反例
            print(f"初始区域{i}没有反例")
        else:
            # 存在反例
            print(f"初始区域{i}反例是 : {counter_ex}")
            cex_init.append(np.array(counter_ex))
            res_init.append(maxi)
    return cex_init, res_init


def MILP_opt_thirdCon(Bx, x_len):
    cex_domain = []
    res_domain = []

    lb_list = [-3, 0]
    ub_list = [0, 2.5]

    lb_domain = domain_min[0]  # -3
    ub_domain = domain_max[0]  # 2.5
    # print("ub: ", ub)
    ly_domain = domain_min[1]  # -2
    uy_domain = domain_max[1]  # 1
    for i in range(len(lb_list)):
        lb = lb_list[i]
        ub = ub_list[i]
        x_l = lb
        x_r = min(round(x_l + x_len, 5), ub)

        while x_l < ub:
            y_l = verify_piecewise_utils.func_x3(x_l)
            y_r = verify_piecewise_utils.func_x3(x_r)
            k, b_u, b_l = verify_piecewise_utils.linear_function_x3(x_l, x_r, y_l, y_r)

            m = gp.Model()

            # 配置
            m.setParam('Outputflag', 0)
            m.setParam('NonConvex', 2)
            # m.setParam('InfUnbdInfo', 1)

            # X 区域约束
            x0 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x0', lb=-np.inf, ub=np.inf)
            x0[0].lb = x_l
            x0[0].ub = x_r

            x0[1].lb = ly_domain
            x0[1].ub = uy_domain

            x1 = m.addMVar((DIM,), vtype=GRB.CONTINUOUS, name='x1', lb=-np.inf, ub=np.inf)
            x1[0].lb = lb_domain
            x1[0].ub = ub_domain

            x1[1].lb = ly_domain
            x1[1].ub = uy_domain

            MILP_theNextPoint(m, x0, x1, k, b_l, b_u)

            y_layer_for_x0 = MILP_Encode_NN(m, x0, Bx, "Bx(x_0)__")
            y_layer_for_x1 = MILP_Encode_NN(m, x1, Bx, "Bx(x_1)__")

            m.setObjective(y_layer_for_x1 - (1 - alpha) * y_layer_for_x0, GRB.MAXIMIZE)

            m.optimize()
            if m.status == 2:
                counter_ex, maxi = solution_output(m)

                if maxi > 0:
                    cex_domain.append(np.array(counter_ex))
                    res_domain.append(maxi)
                    # for v in m.getVars():
                    #     if 'x0' in v.varName or 'x1' in v.varName:
                    #         print(f"domian区域[{x_l},{x_r}]反例是: {v.varName} = {v.x}")

            x_l = x_r
            x_r = min(round(x_l + x_len, 5), ub)
    return cex_domain, res_domain


    # 不存在反例
    #if m.objVal <= gama:
    #    return True, []
    # 存在反例
    #else:
    #    return False, counter_ex


# ——————————————————————————————辅助函数——————————————————————————————

# 将神经网络编码成MILP
def MILP_Encode_NN(m, x0, nn, str = "Bx__"):

    ''' MILP编码神经网络

    :param m: MILP Modle
    :param x0: NN 输入向量 （数学变量）
    :param nn: 需要编码的神经网络
    :param str: Gurobi数学变量 name
    :return: 神经网络最后的输出 （数学变量）
    '''
    W_b_list = ann.serve_for_verify(nn)
    # print("=================================")
    # print("W_b_list: ", W_b_list)
    # print("=================================")
    length = int(len(W_b_list) / 2)

    x_layer = x0
    # x_layer.setAttr(gp.GRB.Attr.VarName, str + "X0")

    for i in range(length):
        W = W_b_list[i * 2]
        b = W_b_list[i * 2 + 1]

        # y = Wx + b
        y_layer = m.addMVar((W.shape[0],), vtype = GRB.CONTINUOUS, lb=-np.inf, ub=np.inf, name = str + f"Y{i + 1}")
        E = np.identity(y_layer.shape[0])
        expr = W @ x_layer + b - E @ y_layer

        name_str = f"layer_{i + 1}, y{i + 1} = W{i + 1}*x{i} + b{i + 1}"
        m.addConstr(expr == 0, name = name_str)

        if i != length - 1:
            # 添加激活函数
            z_layer = m.addMVar((y_layer.shape[0],), vtype=GRB.CONTINUOUS, lb=-np.inf, ub=np.inf, name = str + f"Z{i + 1}")
            for j in range(y_layer.shape[0]):
                m.addConstr(z_layer[j] == gp.max_(y_layer[j], 0.0))
            x_layer = z_layer
        else:
            x_layer = y_layer

        m.update()

    return x_layer


def MILP_theNextPoint(m, x_pre, x_next, k, b_l, b_u):

    b = m.addVar(lb=b_l, ub=b_u, name='b')
    m.addConstr(x_next[0] == x_pre[0] + time_step * x_pre[1])
    m.addConstr(x_next[1] == x_pre[1] + time_step * (- x_pre[0] - x_pre[1] + (k * x_pre[0] + b) / 3.0))

    m.update()

def solution_output(m):
    # print("———————————————————————————————————")
    counter_ex = []
    for v in m.getVars():
        if 'x0' in v.varName:
            counter_ex.append(v.x)
            # print(f"{v.varName} = {v.x}")

        # if 'Bx' in v.varName:
        #    print(f"{v.varName} = {v.x}")

    # print("返回优化结果：")
    # print(m.objVal)

    return counter_ex, m.objVal
'''

# 反例添加
def Counter_Ex_Add(filename, data, flag):
    sampleResult = []
    if flag == 0:
        sampleResult = Sample_Handler.Xi_dataSampling_Near_CounterEx(data)
    elif flag == 1:
        sampleResult = Sample_Handler.Xu_dataSampling_Near_CounterEx(data)
    else:
        sampleResult = Sample_Handler.X_dataSampling_Near_CounterEx(data)

    writeToCsv(filename, sampleResult)

def writeToCsv(filename, data):
    if len(data) == 0:
        return
    data = np.array(data).reshape((-1, x_var_num))
    num = 0
    with open(filename, 'a+', newline='') as f:
        csv_write = csv.writer(f)
        for i in range(data.shape[0]):
            num = num + 1
            csv_write.writerow(data[i, :])
        print(f"{filename},更新了{num}条记录")




def Xi_counter_ex_T_F(Bx, Col, counter_ex, fx_):
    counter_ex = torch.tensor([counter_ex]).to(device)
    Xi_k_point = LossHandler.calc_K_iteration(counter_ex, fx_, Col)
    loss = LossHandler.Xi_Loss_Func(Xi_k_point, Bx)
    if loss > 0:
        #确实是反例
        return False
    else:
        return True

def Xu_counter_ex_T_F(Bx, counter_ex):
    counter_ex = torch.tensor([counter_ex]).to(device)
    pre_y = Bx(counter_ex)
    loss = LossHandler.Xu_Loss_Func(pre_y)
    if loss > 0:
        #确实是反例
        return False
    else:
        return True

def X_counter_ex_T_F(Bx, Col, counter_ex, fx_):
    counter_ex = torch.tensor([counter_ex]).to(device)
    # 第三个蕴含的条件
    pre_z = LossHandler.Filter_Of_Loss3(counter_ex, fx_, Bx, Col)

    loss = 0.
    # 如果没有满足蕴含条件的点
    #print(f"满足蕴含条件的点的个数： {pre_z.shape[0]}")
    if pre_z.shape[0] == 0:
        loss = 0
    else:
        loss = LossHandler.X_Loss_Func(pre_z)

    if loss > 0:
        #确实是反例
        return False
    else:
        return True

'''